{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2018 Google LLC  \n",
    "  \n",
    " Licensed under the Apache License, Version 2.0 (the \"License\");  \n",
    " you may not use this file except in compliance with the License.  \n",
    " You may obtain a copy of the License at  \n",
    "  \n",
    "     http://www.apache.org/licenses/LICENSE-2.0  \n",
    "  \n",
    " Unless required by applicable law or agreed to in writing, software  \n",
    " distributed under the License is distributed on an \"AS IS\" BASIS,  \n",
    " WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  \n",
    " See the License for the specific language governing permissions and  \n",
    " limitations under the License."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Only for when running on Colab:\n",
    "import sys\n",
    "if 'google.colab' in sys.modules:\n",
    "    # Get the dependency .py files, if any.\n",
    "    ! git clone https://github.com/GoogleCloudPlatform/cloudml-samples.git\n",
    "    ! cp cloudml-samples/tpu/templates/tpu_grl_estimator/* .\n",
    "\n",
    "    # Authenticate the user for better GCS access.\n",
    "    # Copy verification code into the text field to continue.\n",
    "    from google.colab import auth\n",
    "    auth.authenticate_user()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow.contrib import summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_CLASSES = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gradient Reversal Layer  \n",
    "  \n",
    " When applied to a tensor this layer is the identity map, but it reverses  \n",
    " the sign of the gradient, and optionally multiplies the reversed gradient  \n",
    " by a weight.  \n",
    "  \n",
    " For details, see [Domain-Adversarial Training of Neural Networks](https://arxiv.org/abs/1505.07818)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GradientReversalLayer(tf.layers.Layer):\n",
    "    def __init__(self, weight=1.0):\n",
    "        super(GradientReversalLayer, self).__init__()\n",
    "        self.weight = weight\n",
    "\n",
    "    def call(self, input_):\n",
    "        @tf.custom_gradient\n",
    "        def _call(input_):\n",
    "            def reversed_gradient(output_grads):\n",
    "                return self.weight * tf.negative(output_grads)\n",
    "\n",
    "            return input_, reversed_gradient\n",
    "        \n",
    "        return _call(input_)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The model function  \n",
    " The network consists of 3 sub-networks:  \n",
    "  \n",
    " * Feature extractor: extracts internal representation for both the source and target distributions.  \n",
    "  \n",
    " * Label predictor: predicts label from the extracted features.  \n",
    "  \n",
    " * Domain classifier: classifies the origin (`source` or `target`) of the extracted features.  \n",
    "  \n",
    "  \n",
    " Both the label predictor and the domain classifier will try to minimize   \n",
    " classification loss, but the gradients backpropagated from the domain  \n",
    " classifier to the feature extractor have their signs reversed.  \n",
    "  \n",
    "  \n",
    " This model function also shows how to use `host_call` to output summaries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_fn(features, labels, mode, params):\n",
    "    source = features['source']\n",
    "    target = features['target']\n",
    "    onehot_labels = tf.one_hot(labels, N_CLASSES)\n",
    "\n",
    "    global_step = tf.train.get_global_step()\n",
    "\n",
    "    # In this sample we use dense layers for each of the sub-networks.\n",
    "    feature_extractor = tf.layers.Dense(7, activation=tf.nn.sigmoid)\n",
    "\n",
    "    label_predictor_logits = tf.layers.Dense(N_CLASSES)\n",
    "\n",
    "    # There are two domains, 0: source and 1: target\n",
    "    domain_classifier_logits = tf.layers.Dense(2)\n",
    "\n",
    "    source_features = feature_extractor(source)\n",
    "    target_features = feature_extractor(target)\n",
    "\n",
    "    # Apply the gradient reversal layer to target features\n",
    "    gr_weight = params['gr_weight']\n",
    "    gradient_reversal = GradientReversalLayer(gr_weight)\n",
    "    target_features = gradient_reversal(target_features)\n",
    "\n",
    "    # The predictions are the predicted labels from the `target` distribution.\n",
    "    predictions = tf.nn.softmax(label_predictor_logits(target_features))\n",
    "    loss = None\n",
    "    train_op = None\n",
    "\n",
    "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "        # define loss\n",
    "        label_prediction_loss = tf.losses.softmax_cross_entropy(\n",
    "            onehot_labels=onehot_labels,\n",
    "            logits=label_predictor_logits(source_features)\n",
    "        )\n",
    "\n",
    "        # There are two domains, 0: source and 1: target\n",
    "        domain_labels = tf.concat((tf.zeros(source.shape[0], dtype=tf.int32), tf.ones(target.shape[0], dtype=tf.int32)), axis=0)\n",
    "        domain_onehot_labels = tf.one_hot(domain_labels, 2)\n",
    "\n",
    "        source_target_features = tf.concat([source_features, target_features], axis=0)\n",
    "\n",
    "        domain_classification_loss = tf.losses.softmax_cross_entropy(\n",
    "            onehot_labels=domain_onehot_labels,\n",
    "            logits=domain_classifier_logits(source_target_features)\n",
    "        )\n",
    "\n",
    "        lambda_ = params['lambda']\n",
    "\n",
    "        loss = label_prediction_loss + lambda_ * domain_classification_loss\n",
    "\n",
    "        # define train_op\n",
    "        optimizer = tf.train.RMSPropOptimizer(learning_rate=0.05)\n",
    "\n",
    "        # wrapper to make the optimizer work with TPUs\n",
    "        if params['use_tpu']:\n",
    "            optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)\n",
    "\n",
    "        train_op = optimizer.minimize(loss, global_step=global_step)\n",
    "\n",
    "    if params['use_tpu']:\n",
    "        # Use host_call to log the losses on the CPU\n",
    "        def host_call_fn(gs, lpl, dcl, ls):\n",
    "            gs = gs[0]\n",
    "            with summary.create_file_writer(params['model_dir'], max_queue=params['save_checkpoints_steps']).as_default():\n",
    "                with summary.always_record_summaries():\n",
    "                    summary.scalar('label_prediction_loss', lpl[0], step=gs)\n",
    "                    summary.scalar('domain_classification_loss', dcl[0], step=gs)\n",
    "                    summary.scalar('loss', ls[0], step=gs)\n",
    "\n",
    "            return summary.all_summary_ops()\n",
    "\n",
    "        # host_call's arguments must be at least 1D\n",
    "        gs_t = tf.reshape(global_step, [1])\n",
    "        lpl_t = tf.reshape(label_prediction_loss, [1])\n",
    "        dcl_t = tf.reshape(domain_classification_loss, [1])\n",
    "        ls_t = tf.reshape(loss, [1])\n",
    "\n",
    "        host_call = (host_call_fn, [gs_t, lpl_t, dcl_t, ls_t])\n",
    "\n",
    "        # TPU version of EstimatorSpec\n",
    "        return tf.contrib.tpu.TPUEstimatorSpec(\n",
    "            mode=mode,\n",
    "            predictions=predictions,\n",
    "            loss=loss,\n",
    "            train_op=train_op,\n",
    "            host_call=host_call)\n",
    "    else:\n",
    "        return tf.estimator.EstimatorSpec(\n",
    "            mode=mode,\n",
    "            predictions=predictions,\n",
    "            loss=loss,\n",
    "            train_op=train_op)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The input function  \n",
    " There are two input data sets, `source` is labeled and `target` is unlabeled."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_input_fn(params={}):\n",
    "    # source distribution: labeled data\n",
    "    source = np.random.rand(100, 5)\n",
    "    labels = np.random.randint(0, N_CLASSES, 100)\n",
    "\n",
    "    # target distribution: unlabeled data\n",
    "    target = np.random.rand(100, 5)\n",
    "\n",
    "    source_tensor = tf.constant(source, dtype=tf.float32)\n",
    "    labels_tensor = tf.constant(labels, dtype=tf.int32)\n",
    "    target_tensor = tf.constant(target, dtype=tf.float32)\n",
    "\n",
    "    # shuffle source and target separately\n",
    "    source_labels_dataset = tf.data.Dataset.from_tensor_slices((source_tensor, labels_tensor)).repeat().shuffle(32)\n",
    "    target_dataset = tf.data.Dataset.from_tensor_slices(target_tensor).repeat().shuffle(32)\n",
    "\n",
    "    # zip them together to set shapes\n",
    "    dataset = tf.data.Dataset.zip((source_labels_dataset, target_dataset))\n",
    "\n",
    "    # TPUEstimator passes params when calling input_fn\n",
    "    batch_size = params.get('batch_size', 16)\n",
    "\n",
    "    dataset = dataset.batch(batch_size, drop_remainder=True)\n",
    "\n",
    "    # TPUs need to know all dimensions when the graph is built\n",
    "    # Datasets know the batch size only when the graph is run\n",
    "    def set_shapes_and_format(source_labels, target):\n",
    "        source, labels = source_labels\n",
    "\n",
    "        source_shape = source.get_shape().merge_with([batch_size, None])\n",
    "        labels_shape = labels.get_shape().merge_with([batch_size])\n",
    "        target_shape = target.get_shape().merge_with([batch_size, None])\n",
    "\n",
    "        source.set_shape(source_shape)\n",
    "        labels.set_shape(labels_shape)\n",
    "        target.set_shape(target_shape)\n",
    "\n",
    "        # Also format the dataset with a dict for features\n",
    "        features = {'source': source, 'target': target}\n",
    "\n",
    "        return features, labels\n",
    "\n",
    "    dataset = dataset.map(set_shapes_and_format)\n",
    "    dataset = dataset.prefetch(tf.contrib.data.AUTOTUNE)\n",
    "\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main(args):\n",
    "    # pass the args as params so the model_fn can use\n",
    "    # the TPU specific args\n",
    "    params = vars(args)\n",
    "\n",
    "    if args.use_tpu:\n",
    "        # additional configs required for using TPUs\n",
    "        tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(args.tpu)\n",
    "        tpu_config = tf.contrib.tpu.TPUConfig(\n",
    "            num_shards=8, # using Cloud TPU v2-8\n",
    "            iterations_per_loop=args.save_checkpoints_steps)\n",
    "\n",
    "        # use the TPU version of RunConfig\n",
    "        config = tf.contrib.tpu.RunConfig(\n",
    "            cluster=tpu_cluster_resolver,\n",
    "            model_dir=args.model_dir,\n",
    "            tpu_config=tpu_config,\n",
    "            save_checkpoints_steps=args.save_checkpoints_steps,\n",
    "            save_summary_steps=100)\n",
    "\n",
    "        # TPUEstimator\n",
    "        estimator = tf.contrib.tpu.TPUEstimator(\n",
    "            model_fn=model_fn,\n",
    "            config=config,\n",
    "            params=params,\n",
    "            train_batch_size=args.train_batch_size,\n",
    "            eval_batch_size=32,\n",
    "            export_to_tpu=False)\n",
    "    else:\n",
    "        config = tf.estimator.RunConfig(model_dir=args.model_dir)\n",
    "\n",
    "        estimator = tf.estimator.Estimator(\n",
    "            model_fn,\n",
    "            config=config,\n",
    "            params=params)\n",
    "\n",
    "    estimator.train(train_input_fn, max_steps=args.max_steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "\n",
    "parser.add_argument(\n",
    "    '--model-dir',\n",
    "    type=str,\n",
    "    default='/tmp/tpu-template',\n",
    "    help='Location to write checkpoints and summaries to.  Must be a GCS URI when using Cloud TPU.')\n",
    "parser.add_argument(\n",
    "    '--max-steps',\n",
    "    type=int,\n",
    "    default=1000,\n",
    "    help='The total number of steps to train the model.')\n",
    "parser.add_argument(\n",
    "    '--train-batch-size',\n",
    "    type=int,\n",
    "    default=16,\n",
    "    help='The training batch size.  The training batch is divided evenly across the TPU cores.')\n",
    "parser.add_argument(\n",
    "    '--save-checkpoints-steps',\n",
    "    type=int,\n",
    "    default=100,\n",
    "    help='The number of training steps before saving each checkpoint.')\n",
    "parser.add_argument(\n",
    "    '--use-tpu',\n",
    "    action='store_true',\n",
    "    help='Whether to use TPU.')\n",
    "parser.add_argument(\n",
    "    '--tpu',\n",
    "    default=None,\n",
    "    help='The name or GRPC URL of the TPU node.  Leave it as `None` when training on AI Platform.')\n",
    "\n",
    "parser.add_argument(\n",
    "    '--gr-weight',\n",
    "    default=1.0,\n",
    "    help='The weight used in the gradient reversal layer.')\n",
    "parser.add_argument(\n",
    "    '--lambda',\n",
    "    default=1.0,\n",
    "    help='The trade-off between label_prediction_loss and domain_classification_loss.')\n",
    "\n",
    "args, _ = parser.parse_known_args()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# colab.research.google.com specific\n",
    "if 'google.colab' in sys.modules:\n",
    "    import json\n",
    "    import os\n",
    "\n",
    "    # TODO(user): change this\n",
    "    args.model_dir = 'gs://your-gcs-bucket'\n",
    "\n",
    "    # When connected to the TPU runtime\n",
    "    if 'COLAB_TPU_ADDR' in os.environ:\n",
    "        tpu_grpc = 'grpc://{}'.format(os.environ['COLAB_TPU_ADDR'])\n",
    "\n",
    "        args.tpu = tpu_grpc\n",
    "        args.use_tpu = True\n",
    "\n",
    "        # Upload credentials to the TPU\n",
    "        with tf.Session(tpu_grpc) as sess:\n",
    "            data = json.load(open('/content/adc.json'))\n",
    "            tf.contrib.cloud.configure_gcs(sess, credentials=data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "main(args)"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
